#include <dialects/CFG/CFG.h>
#include <dialects/Func/Func.h>
#include <dialects/PL/PL.h>
#include <dialects/PL/Translation/ToCFG.h>
#include <lcir/IR/Expression.h>
#include <lcir/IR/IRBuilder.h>
#include <lcir/Pass/Pass.h>
#include <dialects/CFG/Ops.h>
#include <dialects/CFG/TyAndProto.h>
#include <dialects/PL/Ops.h>
#include <dialects/PL/TypesAndProtos.h>
#include <lcir/IR/Argument.h>
#include <lcir/IR/ExprTemps.h>
#include <lcir/IR/Lambda.h>

#include <cassert>
#include <iostream>
#include <vector>
namespace lcir::pl {
/*
if () {""}
else if {}

handle () {
        handleCond(cond, ifBlock, elBlock) {
                handle(lhs, )
                handle(rhs)
        }
        ifblock
        elblock
}


*/

static bool isSealed(Value* val) {
  auto* blockProto = val->getProto<cfg::BlockProto>();
  assert(blockProto && "val is not a block");
  return blockProto->isSealed();
}
static bool isSealed(Lambda* val) {
  auto* blockProto = val->getProto<cfg::BlockProto>();
  assert(blockProto && "val is not a block");
  return blockProto->isSealed();
}

static Value* handleCond(Value* cond, Value* trueBlock, Value* falseBlock, IRBuilder& b) {
  auto* expr = cond->getDefExpr();
  if (!expr || (!expr->isOp<And>() && !expr->isOp<Or>())) {
    auto* blockLambda = Lambda::get();
    auto* block = b.createAndGetRet<cfg::Block>({blockLambda});
    block->setInfo(expr->getOp()->getName());

    auto insertPoint = b.getInsertPoint();
    b.setInsertPoint(blockLambda);
    b.createAndGetRet<cfg::CondBr>({cond, trueBlock, falseBlock});
    b.setInsertPoint(insertPoint);
    return block;
  }
  auto insertPoint = b.getInsertPoint();
  // b.setInsertPoint(expr);
  // 给And/OR创建一个block
  Value *lhsBlock{nullptr}, *rhsBlock{nullptr};
  rhsBlock = handleCond(expr->getOperandValue(1), trueBlock, falseBlock, b);
  b.setInsertPoint(rhsBlock->getDefExpr());
  if (expr->isOp<And>()) {
    lhsBlock = handleCond(expr->getOperandValue(0), rhsBlock, falseBlock, b);
  } else if (expr->isOp<Or>()) {
    lhsBlock = handleCond(expr->getOperandValue(0), trueBlock, rhsBlock, b);
  }
  auto* lhsLambda = lhsBlock->getProto<cfg::BlockProto>()->getLambda();
  auto* rhsLambda = rhsBlock->getProto<cfg::BlockProto>()->getLambda();

  b.setInsertPoint(insertPoint);
  return lhsBlock;
}

static void handleIf(Expression* expr, Value* currentBlock, IRBuilder& b) {
  IfTemp temp(expr);
  auto* cond = temp.getCond();
  auto* ifLmabda = Lambda::get();
  // auto* ifLmabda = temp.getLambda()->makeNewLambda();
  auto* ifBlock = b.createAndGetRet<cfg::Block>({ifLmabda});
  // 看看他后面有没有跟else
  if (auto* condBlock = expr->getResult()->getProto()->cast<pl::CondBlockProto>()) {
    if (auto* elCondBlock = condBlock->getNext()) {
      auto* el = elCondBlock->getValue()->getDefExpr();
      auto* elLambda = Lambda::get();
      auto* elBlock = b.createAndGetRet<cfg::Block>({elLambda});
    }
  }

  auto* elLambda = Lambda::get();
  auto* elBlock = b.createAndGetRet<cfg::Block>({elLambda});
  ifBlock->setInfo("if");
  elBlock->setInfo("else");
  // b.setInsertPoint(ifBlock->getDefExpr());

  // 处理逻辑短路
  IRBuilder condb(ifBlock->getDefExpr());
  auto* condBlock = handleCond(cond, ifBlock, elBlock, condb);
  b.setInsertPoint(currentBlock->getProto<cfg::BlockProto>()->getLambda());
  b.create<cfg::Br>({condBlock});
  b.setInsertPoint(expr);
  b.createAndGetRet<cfg::CondBr>({cond, ifBlock, elBlock});
}
/**
 * @param lambda 正在处理的原始lambda
 * @param currentBlockLambda 新的lambda
 */
static Lambda* handleLambda(Lambda* lambda, Lambda* currentBlockLambda, Value* loopCond, Value* loopEnd) {
  // auto* currentBlock = b.createAndGetRet<cfg::Block>({currentBlockLambda});
  for (auto& expr : IncRange(lambda->getExprs())) {
    // b的插入点 新lambda的父lambda, 用来createBlock的
    IRBuilder b(currentBlockLambda->getDefExpr()->getLambda());

    // b1的插入点 直接在新lambda里面插入东西
    IRBuilder b1(currentBlockLambda);
    if (expr.isOp<If>()) {
      // IRBuilder b(currentBlockLambda->getDefExpr()->getLambda());
      IfTemp temp(&expr);
      // 创建if block和 end Block
      auto* ifBlock = b.createAndGetRet<cfg::Block>({Lambda::get()});
      ifBlock->setInfo("if");

      // 这里看下是if还是if-else
      Value* elBlock{nullptr};
      if (auto* condBlock = expr.getResult()->getProto()->cast<pl::CondBlockProto>()) {
        if (auto* elCondBlock = condBlock->getNext()) {
          elBlock = b.createAndGetRet<cfg::Block>({Lambda::get()});
          elBlock->setInfo("else");
        }
      }
      auto* endBlock = b.createAndGetRet<cfg::Block>({Lambda::get()});
      endBlock->setInfo("end");

      IfTemp ifTemp(&expr);
      IRBuilder condb(ifBlock->getDefExpr());
      auto cond = handleCond(temp.getCond(), ifBlock, elBlock ? elBlock : endBlock, condb);

      IRBuilder b1(currentBlockLambda);
      b1.create<cfg::Br>({cond});

      auto* ifLambdaBlock =
          handleLambda(ifTemp.getLambda(), ifBlock->getProto<cfg::BlockProto>()->getLambda(), loopCond, loopEnd);
      if (elBlock) {
        auto* elLambdaBlock =
            handleLambda(expr.getResult()->getProto()->cast<pl::CondBlockProto>()->getNext()->getLambda(),
                         elBlock->getProto<cfg::BlockProto>()->getLambda(), loopCond, loopEnd);
        b1.setInsertPoint(elLambdaBlock);
        b1.create<cfg::Br>({endBlock});
      }
      b1.setInsertPoint(ifLambdaBlock);
      b1.create<cfg::Br>({endBlock});
      currentBlockLambda = endBlock->getProto<cfg::BlockProto>()->getLambda();
    } else if (expr.isOp<El>()) {
      continue;
    } else if (expr.isOp<Loop>()) {
      // IRBuilder b(currentBlockLambda->getDefExpr()->getLambda());
      LoopTemp temp(&expr);
      auto* loopBlock = b.createAndGetRet<cfg::Block>({Lambda::get()});
      auto* endBlock = b.createAndGetRet<cfg::Block>({Lambda::get()});
      endBlock->setInfo("loopend");
      IRBuilder condb(loopBlock->getDefExpr());
      auto condBlock = handleCond(temp.getCond(), loopBlock, endBlock, condb);
      std::cout << "currentBlockLambda" << "\n";
      
      if (!currentBlockLambda->isSealed()) b1.create<cfg::Br>({condBlock});

      auto* loopLambdaBlock =
          handleLambda(temp.getLambda(), loopBlock->getProto<cfg::BlockProto>()->getLambda(), condBlock, endBlock);
      loopLambdaBlock->getDefExpr()->getResult()->setInfo("loopBody");
      b1.setInsertPoint(loopLambdaBlock);
      if (!isSealed(condBlock)) b1.create<cfg::Br>({condBlock});
      currentBlockLambda = endBlock->getProto<cfg::BlockProto>()->getLambda();
    } else if (expr.isOp<pl::Break>()) {
      assert(loopEnd && "break outside loop");
      if (!currentBlockLambda->isSealed()) b1.create<cfg::Br>({loopEnd});
      break;
    } else if (expr.isOp<pl::Continue>()) {
      assert(loopCond && "continue outside loop");
      if (!currentBlockLambda->isSealed()) b1.create<cfg::Br>({loopCond});
      break;
    } else {
      auto uptr = expr.removeFromLambda();
      currentBlockLambda->pushBackExpr(std::move(uptr));
    }
  }
  return currentBlockLambda;
}

static Lambda* handleFunc(Lambda* func) {
  // 搞一个新lambda
  auto* newLambda = Lambda::get();
  auto lambdaArgs = func->getArgs();
  for(auto [idx, arg]: IndexRange(lambdaArgs)) {
    newLambda->addArg(arg->getType());
    arg->rauw(newLambda->getArg(idx));
  }
  IRBuilder b(newLambda);
  auto* currentBlockLambda = Lambda::get();
  auto* currentBlock = b.createAndGetRet<cfg::Block>({currentBlockLambda});
  currentBlock->setInfo("entry");
  handleLambda(func, currentBlockLambda, nullptr, nullptr);
  return newLambda;
}

void ToCFG::run() {
  auto* module = getLambda();
  std::vector<Expression*> funcs;
  module->walkOnOp<func::Func>([&](Expression* expr) { funcs.push_back(expr); });
  for (auto* func : funcs) {
    func::FuncTemp temp(func);
    auto* lambda = handleFunc(temp.getFuncBody());
    func->getOperandValue(0)->rauw(lambda);
  }
  // entry必然是第一个
  module->walkOnOp<func::Func>([&](Expression* expr) {
    func::FuncTemp temp(expr);
    auto* entry = temp.getFuncBody()->begin()->getResultProto()->dyn_cast<cfg::BlockProto>();
    auto blocks = cfg::reversePostorderTraversal(entry);
    // for(auto block: blocks) {
    //   std::cout << block->getInfo() << "\n";
    // }
  });
}
}  // namespace lcir::pl